import numpy as np
import h5py
from utils.config import *
from utils.train import train_model
from utils.evaluate import evaluate_model
from utils.evaluate import plot_forecast_segments, calculate_rmse_BCD
from utils.data_loader import process_images
#其他参数配置请查看utils/config文件
#可选择参数:["MLP", "LSTM", "CNN_LSTM"]
config.MODEL_SELECT = "LSTM"

if __name__ == "__main__":
    """ 
    光伏预测
    """
    # 设置随机种子
    set_random_seed()
    # 配置GPU，没有GPU将用CPU训练
    configure_hardware()
    #weathers = ["CLOUDY", "SUNNY", "OVERCAST","ALL"]
    weathers = ["ALL"]
    metrics_list = []
    weather_list = []
    mae_list = []
    for weather_select in weathers:
        # 配置天气模式
        config.WEATHER_SELECT = weather_select
        # 读取数据集训练
        data_path = os.path.join(config.data_folder, "forecast.hdf5")
        times_trainval = np.load(os.path.join(config.data_folder, "times_trainval.npy"), allow_pickle=True)
        times_test = np.load(os.path.join(config.data_folder, "times_test.npy"), allow_pickle=True)
        # 训练模型
        history, model = train_model(data_path, times_trainval)
        # 读取测试集
        with h5py.File(data_path, 'r') as f:
            test_images = process_images(f['test']['image_log_test'][:,-1,:,:,:])#(None, 16, 64, 64, 3)
            test_pv_log = f['test']['pv_log_test'][...].astype('float32')
            test_pv_pred = f['test']['pv_pred_test'][...].astype('float32')
        # 模型测试与评估
        test_pv_pred = test_pv_pred[:,-1]
        metrics, predictions = evaluate_model(model, (test_images, test_pv_log, test_pv_pred), times_test)
        for metric_name, metric_value in metrics['overall'].items():
            print(f"{weather_select} {metric_name}: {metric_value:.3f}")

        metrics_list.append(metrics['overall']['RMSE'])
        mae_list.append(metrics['overall']['MAE'])
        weather_list.append(len(predictions))
        times = [np.datetime64(t) for t in times_test]

        plot_forecast_segments(weather_select, times, predictions, test_pv_pred)
    #rmse_BCD, rmse_A = calculate_rmse_BCD(metrics_list, mae_list,weather_list)


    
        
            
            
            